package com.dtdream.emr;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorConverters;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.io.Text;

import java.util.ArrayList;

/**
 * Created by moriarty on 17-8-4.
 */
@Description("_FUNC_(col, separator) - split column into several rows by separator")
public class SplitWordsUdf extends GenericUDF {

    private transient ObjectInspectorConverters.Converter[] converters;


    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        if (arguments.length != 2) {
            throw new UDFArgumentException("Function GenericUDTFSplit requires 2 arguments, got " + arguments.length);
        }

        this.converters = new ObjectInspectorConverters.Converter[arguments.length];

        for(int i = 0; i < arguments.length; ++i) {
            this.converters[i] = ObjectInspectorConverters.getConverter(arguments[i], PrimitiveObjectInspectorFactory.writableStringObjectInspector);
        }

        return ObjectInspectorFactory.getStandardListObjectInspector(PrimitiveObjectInspectorFactory.writableStringObjectInspector);
    }

    public Object evaluate(GenericUDF.DeferredObject[] arguments) throws HiveException {
        assert arguments.length == 2;

        if (arguments[0].get() != null && arguments[1].get() != null) {
            Text s = (Text)this.converters[0].convert(arguments[0].get());
            Text regex = (Text)this.converters[1].convert(arguments[1].get());
            ArrayList<Text> result = new ArrayList();
            String[] arr$ = s.toString().replaceAll("[`~!@#$^&*()=|{}\\-':;',\\[\\].<>/?~！@#￥……&*（）——|{}【】‘；：”“'。，、？]|(\\s+(a|an|the)\\s+)", regex.toString()).split(regex.toString());
            int len$ = arr$.length;

            for(int i$ = 0; i$ < len$; ++i$) {
                String str = arr$[i$];
                if (str == null || "".equals(str) || str.length() == 1) {
                    continue;
                }
                result.add(new Text(str));
            }

            return result;
        } else {
            return null;
        }
    }

    public String getDisplayString(String[] children) {
        assert children.length == 2;

        return this.getStandardDisplayString("splitword", children);
    }
}
